﻿using System.Collections.Generic;
using System.Linq;
using System.Text;
using System;
using System.Diagnostics;
using Microsoft.Diagnostics.Runtime;
using System.Threading;
using System.IO;

namespace WMemoryProfiler.Memory
{
    /// <summary>
    /// Can dump memory statistics from your own or external processes. 
    /// </summary>
    public class ProcessMemoryDumper : IDisposable
    {
        int ProcessId;
        bool UseMemoryDumps;
        string MemoryDumpFile;

        /// <summary>
        /// Create a memory dumpr for a process with id with default settings (do close debugger after every dump and do hide debugger window)
        /// </summary>
        /// <param name="processId"></param>
        public ProcessMemoryDumper(int processId):this(processId, false)
        {
        }

        /// <summary>
        /// Create a memory dumper for a process with id with specific settings. 
        /// </summary>
        /// <param name="processId">process id to watch</param>
        /// <param name="useMemoryDumps">When true instead of attaching ot a process a memory dump is taken and analyzed. This can be useful
        /// for later post mortem analysis and it also mitigates some debugger hangs which were not solvable at all.</param>
        public ProcessMemoryDumper(int processId, bool useMemoryDumps)
        {
            ProcessId = processId;
            UseMemoryDumps = useMemoryDumps;
        }

        /// <summary>
        /// Open an existing memory dump file
        /// </summary>
        /// <param name="dumpFile">Existing dump file</param>
        public ProcessMemoryDumper(string dumpFile)
        {
            UseMemoryDumps = true;
            if ( String.IsNullOrEmpty(dumpFile))
            {
                throw new ArgumentException("dumpFile");
            }

            if( !File.Exists(dumpFile) )
            {
                throw new FileNotFoundException(String.Format("Memory dump file {0} was not found!", dumpFile));
            }

            MemoryDumpFile = dumpFile;
        }

        /// <summary>
        /// Get the memory consumption per type summarized.
        /// </summary>
        /// <returns></returns>
        public MemoryStatistics GetMemoryStatistics()
        {
            DataTarget target = null;
DumpAgain:
            try
            {
                if (UseMemoryDumps)
                {
                    if (MemoryDumpFile == null)
                    {
                        var dumper = new ProcessDumpCreator("ForMemoryProfiling");
                        MemoryDumpFile = dumper.Dump(this.ProcessId);
                        if (MemoryDumpFile == null)
                        {
                            throw new InvalidOperationException(String.Format("Could not create memory dump from process with id: {0}.", this.ProcessId));
                        }
                    }

                    target = DataTarget.LoadCrashDump(MemoryDumpFile);
                }
                else
                {
                    target = DataTarget.AttachToProcess(this.ProcessId, 5000);
                }

                if( target.ClrVersions.Count == 0 )
                {
                    throw new InvalidOperationException(String.Format("No CLR Runtime found in process with id: {0}", this.ProcessId));
                }

                ClrRuntime runtime = target.ClrVersions[0].CreateRuntime();

                ClrHeap heap = runtime.GetHeap();
                if (!heap.CanWalkHeap) // wait for gc to complete
                {
                    goto DumpAgain;
                }

                ObjectSet liveObjects = GetLiveObjects(heap); // make sure that we get only rooted objects to prevent reporting wrong instance counts of already dead objects

                var stats = from o in heap.EnumerateObjectAddresses()
                            let t = heap.GetObjectType(o)
                            where t != null && liveObjects.Contains(o)
                            group o by t into g
                            let size = g.Sum(o => (uint)g.Key.GetSize(o))
                            orderby size
                            select new
                            {
                                Name = g.Key.Name,
                                Size = size,
                                Count = g.Count()
                            };

                Dictionary<string, InstanceInfo> instanceInfos = new Dictionary<string, InstanceInfo>();

                foreach (var stat in stats)
                {
                    InstanceInfo info = null;
                    if (!instanceInfos.TryGetValue(stat.Name, out info))
                    {
                        info = new InstanceInfo();
                        instanceInfos.Add(stat.Name, info);
                    }

                    info.InstanceCount += stat.Count;
                    info.AllocatedBytes += stat.Size;
                }

                using (var proc = Process.GetProcessById(this.ProcessId))
                {
                    var memStat = new MemoryStatistics(instanceInfos, proc.PrivateMemorySize64, proc.WorkingSet64, proc.VirtualMemorySize64);
                    return memStat;
                }
            }
            finally
            {
                if( target != null )
                {
                    target.Dispose();
                }
            }

        }

        /// <summary>
        /// Taken from https://harshaprojects.wordpress.com/2015/12/29/clr-md-analyzing-live-process/
        /// </summary>
        /// <param name="heap">Clr Heap</param>
        /// <returns>ObjectSet with all live object addresses.</returns>
        private static ObjectSet GetLiveObjects(ClrHeap heap)
        {
            ObjectSet considered = new ObjectSet(heap);
            Stack<ulong> eval = new Stack<ulong>();

            foreach (var root in heap.EnumerateRoots())
                eval.Push(root.Object);

            while (eval.Count > 0)
            {
                ulong obj = eval.Pop();
                if (considered.Contains(obj))
                    continue;

                considered.Add(obj);

                var type = heap.GetObjectType(obj);
                if (type == null)  // Only if heap corruption
                    continue;

                type.EnumerateRefsOfObjectCarefully(obj, delegate (ulong child, int offset)
                {
                    if (child != 0 && !considered.Contains(child))
                        eval.Push(child);
                });
            }

            return considered;
        }


        /// <summary>
        /// Generate from two memory statistics a diff to see what has changed on the managed heap
        /// </summary>
        /// <param name="old">"old" memory snapshot</param>
        /// <param name="later">"new" memory snapshot</param>
        /// <returns>Array of diff values sorted by diff number in descending order</returns>
        public InstanceInfoDiff[] GetMemoryStatisticsDiff(MemoryStatistics old, MemoryStatistics later)
        {
            return MemoryDiff(old.ManagedHeapStats, later.ManagedHeapStats);
        }

        /// <summary>
        /// Get current memory stats and compare it to a previous version
        /// </summary>
        /// <param name="old">old stats</param>
        /// <returns>Array of diff values sorted by diff number in descending order</returns>
        public InstanceInfoDiff[] GetMemoryStatisticsDiff(MemoryStatistics old)
        {
            return GetMemoryStatisticsDiff(old, GetMemoryStatistics());
        }

        /// <summary>
        /// Generate from two memory statistics a diff to see what has changed on the managed heap
        /// </summary>
        /// <param name="old">"old" memory snapshot</param>
        /// <param name="later">"new" memory snapshot</param>
        /// <returns>Array of diff values sorted by diff number in descending order</returns>
        public InstanceInfoDiff[] MemoryDiff(Dictionary<string, InstanceInfo> old, Dictionary<string, InstanceInfo> later)
        {
            List<InstanceInfoDiff> lret = new List<InstanceInfoDiff>();
            HashSet<string> addedRemvoed = new HashSet<string>();

            // types which do not occur anymore in memory
            foreach (var removedType in old.Keys.Where(oldType => !later.ContainsKey(oldType)))
            {
                addedRemvoed.Add(removedType);
                var oldInfo = old[removedType];
                lret.Add(new InstanceInfoDiff(0, oldInfo.InstanceCount, 0, oldInfo.AllocatedBytes, removedType));
            }

            // types which are new in memory
            foreach (var addedType in later.Keys.Where(newType => !old.ContainsKey(newType)))
            {
                addedRemvoed.Add(addedType);
                var added = later[addedType];
                lret.Add(new InstanceInfoDiff(added.InstanceCount, 0, added.AllocatedBytes, 0, addedType));
            }

            // Types with changes
            foreach (var changed in old.Keys.Where(type => !addedRemvoed.Contains(type)))
            {
                InstanceInfo oldInfo = old[changed];
                InstanceInfo newInfo = later[changed];
                var change = new InstanceInfoDiff(newInfo.InstanceCount, oldInfo.InstanceCount, newInfo.AllocatedBytes, oldInfo.AllocatedBytes, changed);

                lret.Add(change);
            }

            return lret.OrderByDescending(x => x.InstanceCountDiff).ToArray();
        }

        /// <summary>
        /// Close all open debuggers of this instance and other resources
        /// </summary>
        public void Dispose()
        {
            Dispose(true);
        }

        /// <summary>
        /// Dispose for child classes
        /// </summary>
        /// <param name="isDisposing"></param>
        protected virtual void Dispose(bool isDisposing)
        {

        }
    }
}
